import utils as utils
import cv2
import torch
import einops
from model import make_model

torch.nn.Transformer.generate_square_subsequent_mask

model = make_model(10,10,encoder_layers=1,decoder_layers=1,d_model=8,d_ff=32).eval()

src = torch.LongTensor([[0, 1, 2, 3, 4]])
tgt = torch.LongTensor([[4, 3, 2, 1, 0]])

print("no target mask")
# model(src,tgt[:,:1])

mask_1 = utils.get_subsequent_mask(1)
mask_2 = utils.get_subsequent_mask(2)
print("With target mask")

# print(model(src,tgt[:,:1]))
# print(model(src,tgt[:,:1]))
# model(src,tgt[:,:1])
print(model(src,tgt[:,:1],target_mask=mask_1))
print()
print(model(src,tgt[:,:1],target_mask=mask_2))
print()
print(model(src,tgt,target_mask=utils.get_subsequent_mask(5)))


# d = torch.randn([1,3,3])
# mask_d = utils.get_subsequent_mask(3)
# # mask_d.unsqueeze_(-1)
# d.masked_fill_(mask=mask_d,value=float('-inf'))
# print(d)

# mask = torch.nn.Transformer.generate_square_subsequent_mask(2)
# mask2 = utils.get_subsequent_mask(2)
# print(mask)
# print(mask2)